﻿#include "RtmpPushStream.h"

#include <QDebug>
#include <QCoreApplication>
#include <QElapsedTimer>
#include <QJsonDocument>
#include <QJsonObject>

#define __STDC_CONSTANT_MACROS

RtmpPushStream::RtmpPushStream(QObject *parent)
    : QThread(parent)
    , m_bInitSuccess(false)
    , m_bThreadRun(true)
{
    QString filePath = QCoreApplication::applicationDirPath() + "/config.json";
    QFile file(filePath);
    if(file.open(QIODevice::ReadOnly))
    {
        QJsonObject json = QJsonDocument::fromJson(file.readAll()).object();
        m_rtmpAddr = json["addr"].toString();
    }

    m_image_width = 1280;
    m_image_height = 720;

    this->start();
}

RtmpPushStream::~RtmpPushStream()
{

}


AVCodecContext *RtmpPushStream::init_encoder(int width, int height)
{
    // AVCodec 是 FFmpeg 内部全局对象，通过 avcodec_find_encoder 获取的是指针引用，无需手动释放。
    AVCodec *codec = (AVCodec *)avcodec_find_encoder(AV_CODEC_ID_H264);
    if(!codec){
        qDebug() << __FILE__ << __FUNCTION__ << __LINE__ <<"avcodec_find_encoder error";
        return nullptr;
    }

    AVCodecContext *codec_ctx = avcodec_alloc_context3(codec);
    if(!codec_ctx)
    {
        qDebug() << __FILE__ << __FUNCTION__ << __LINE__ <<"avcodec_alloc_context3 error";
        return nullptr;
    }
    codec_ctx->width = width;
    codec_ctx->height = height;
    // // 必须与输入帧格式一致
    codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    // // 帧率分母/分子（如30fps）
    codec_ctx->time_base = AVRational{1, 30};
    codec_ctx->framerate = AVRational{30, 1};
    // // 4 Mbps
    codec_ctx->bit_rate = 4000000;
    // // 关键帧间隔
    codec_ctx->gop_size = 30;
    // // B帧数量
    codec_ctx->max_b_frames = 0;

    // 强制下一个帧为IDR
    av_opt_set(codec_ctx->priv_data, "forced-idr", "1", 0);

    if (avcodec_open2(codec_ctx, codec, nullptr) < 0) {
        qDebug() << __FILE__ << __FUNCTION__ << __LINE__ << "Failed to open encoder";
        // 释放已分配的codec_ctx
        avcodec_free_context(&codec_ctx);
        return nullptr;
    }
    return codec_ctx;
}

AVFormatContext *RtmpPushStream::init_rtmp_output(const char *rtmp_url, AVCodecContext *codec_ctx)
{
    if(m_rtmpAddr.isEmpty())
    {
        qDebug() << "rtmp addr is empty";
        return nullptr;
    }
    AVFormatContext *fmt_ctx = nullptr;
    int ret = 0;
    // 1. 初始化输出上下文
    ret = avformat_alloc_output_context2(&fmt_ctx, nullptr, "flv", rtmp_url);
    if (ret < 0 || !fmt_ctx) {
        qDebug() <<__FILE__ << __FUNCTION__ << __LINE__ << "Failed to allocate output context: " << ret;
        return nullptr;  // 无需释放，fmt_ctx 分配失败时无资源占用
    }

    // 2. 创建输出流
    AVStream *stream = avformat_new_stream(fmt_ctx, nullptr);
    if (!stream) {
        qDebug() <<__FILE__ << __FUNCTION__ << __LINE__ << "Failed to create output stream";
        avformat_free_context(fmt_ctx);  // 释放已分配的 fmt_ctx
        return nullptr;
    }

    // 3. 从编码器复制参数到流
    ret = avcodec_parameters_from_context(stream->codecpar, codec_ctx);
    if (ret < 0) {
        qDebug() <<__FILE__ << __FUNCTION__ << __LINE__ << "Failed to copy codec parameters: " << ret;
        avformat_free_context(fmt_ctx);
        return nullptr;
    }

    // 4. 打开输出（RTMP 服务器）
    if (!(fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&fmt_ctx->pb, rtmp_url, AVIO_FLAG_WRITE);
        if (ret < 0) {
            qDebug() << __FILE__ << __FUNCTION__ << __LINE__ << "Failed to open RTMP URL: " << ret;
            avformat_free_context(fmt_ctx);
            return nullptr;
        }
    }

    // 5. 写入头部信息
    ret = avformat_write_header(fmt_ctx, nullptr);
    if (ret < 0) {
        qDebug() << __FILE__ << __FUNCTION__ << __LINE__ <<"Failed to write header: " << ret;
        // // 关闭已打开的 IO 上下文
        if (fmt_ctx->pb)
            avio_closep(&fmt_ctx->pb);
        avformat_free_context(fmt_ctx);
        return nullptr;
    }

    return fmt_ctx;  // 成功则返回上下文
}

bool RtmpPushStream::update_encoder_and_stream(int image_width, int image_height)
{
    if(h264_codec_ctx)
    {
        m_image_width = image_width;
        m_image_height = image_height;

        // 1. 刷新旧编码器（发送缓存帧）
        avcodec_send_frame(h264_codec_ctx, nullptr);

        // 2. 关闭旧编码器
        avcodec_free_context(&h264_codec_ctx);

        // 3. 创建新编码器
        h264_codec_ctx = init_encoder(image_width, image_height);

        // 4. 更新输出流参数
        avcodec_parameters_from_context(out_stream->codecpar, h264_codec_ctx);

        // 5. 强制下一帧为关键帧
        // key_frame_requested = true;

        avformat_write_header(h264_ofmt_ctx, NULL); // 重新写入头
    }
    else{
        m_image_width = image_width;
        m_image_height = image_height;
    }

    return true;
}

void RtmpPushStream::addPacket(AVPacket *packet)
{
    QMutexLocker lk(&m_packetMutex);
    m_packetList.append(packet);
}

int RtmpPushStream::splitNalus(const uint8_t *data, int size)
{
    int idrIndex = -1;

    const uint8_t* start = data;
    const uint8_t* end = data + size;

    while (start < end) {
        // 查找起始码 (00 00 01 或 00 00 00 01)
        if (start + 3 < end && start[0] == 0 && start[1] == 0 && start[2] == 1) {
            int start_code_len = (start[3] == 1) ? 4 : 3;
            const uint8_t* nalu_start = start + start_code_len;

            // 查找下一个起始码
            const uint8_t* next = nalu_start;
            while (next + 2 < end && !(next[0] == 0 && next[1] == 0 && next[2] == 1)) {
                next++;
            }

            // 提取NAL单元
            //int nalu_size = next - nalu_start;
            uint8_t nal_type = nalu_start[0] & 0x1F;
            //qDebug() << "Found NAL unit: type=" << nal_type << "size=" << nalu_size;

            if(nal_type == 5)
            {
                idrIndex = start - data;
                return idrIndex;
            }

            start = next;

        } else {
            start++;
        }
    }

    return idrIndex;
}

/**
 * @brief create_sei_nal_unit 不处理竞争版本
 * @param sei_text
 * @param sei_text_size
 * @param sei_nalu_size
 * @return
 */
uint8_t *RtmpPushStream::create_sei_nal_unit(const char *sei_text, int sei_text_size, int *sei_nalu_size)
{
    const int payload_type = 5; // user_data_unregistered
    const uint8_t uuid[] = {0xdc,0x45,0xe9,0xbd,0xe6,0xd9,0x48,0xb7,0x96,0x2c,0xd8,0x20,0xd9,0x23,0xee,0xef};

    // 1. 合并UUID、用户数据和结束标记0x80
    uint8_t* combined_payload = (uint8_t*)av_malloc(16 + sei_text_size + 1);
    memcpy(combined_payload, uuid, 16);
    memcpy(combined_payload + 16, sei_text, sei_text_size);
    combined_payload[16 + sei_text_size] = 0x80; // RBSP尾部标记
    int combined_size = 16 + sei_text_size + 1;

    int rbsp_size = combined_size;

    // 3. 计算UVLC编码字节数
    int uvlc_bytes = 1;
    int tmp = rbsp_size;
    while (tmp >= 0xFF) {
        uvlc_bytes++;
        tmp -= 0xFF;
    }

    // 4. 分配内存（起始码4 + NAL头1 + payload_type1 + uvlc + rbsp_data）
    int nal_size = 4 + 1 + 1 + uvlc_bytes + rbsp_size;
    uint8_t* sei_nal = (uint8_t*)av_malloc(nal_size + AV_INPUT_BUFFER_PADDING_SIZE);
    if (!sei_nal) {
        av_free(combined_payload);
        return NULL;
    }
    memset(sei_nal + nal_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);

    // 5. 填充起始码和NAL头
    sei_nal[0] = 0x00; sei_nal[1] = 0x00;
    sei_nal[2] = 0x00; sei_nal[3] = 0x01;
    // SEI NAL type
    sei_nal[4] = 0x06;

    // 6. 填充SEI payload
    uint8_t* p = sei_nal + 5;
    *p++ = payload_type;

    // 6.1 UVLC编码的payload_size
    tmp = rbsp_size;
    while (tmp >= 0xFF) {
        *p++ = 0xFF;
        tmp -= 0xFF;
    }
    *p++ = tmp;

    // 6.2 填充RBSP数据（处理防竞争字节）
    memcpy(p, combined_payload, combined_size);
    av_free(combined_payload);
    p += combined_size;

    // 7. 验证大小
    if ((p - sei_nal) != nal_size) {
        av_log(NULL, AV_LOG_ERROR, "Size mismatch: %td vs %d\n",
               p - sei_nal, nal_size);
        av_free(sei_nal);
        return NULL;
    }

    *sei_nalu_size = nal_size;
    return sei_nal;
}

/**
 * @brief create_sei_nal_unit_emu 处理竞争版本
 * @param sei_text
 * @param sei_text_size
 * @param sei_nalu_size
 * @return
 */
uint8_t *create_sei_nal_unit_emu(const char *sei_text, int sei_text_size, int *sei_nalu_size)
{
    const int payload_type = 5; // user_data_unregistered
    const uint8_t uuid[] = {0xdc,0x45,0xe9,0xbd,0xe6,0xd9,0x48,0xb7,0x96,0x2c,0xd8,0x20,0xd9,0x23,0xee,0xef};

    // 1. 合并UUID、用户数据和结束标记0x80
    uint8_t* combined_payload = (uint8_t*)av_malloc(16 + sei_text_size + 1);
    memcpy(combined_payload, uuid, 16);
    memcpy(combined_payload + 16, sei_text, sei_text_size);
    combined_payload[16 + sei_text_size] = 0x80; // RBSP尾部标记
    int combined_size = 16 + sei_text_size + 1;

    // 2. 计算防竞争字节数
    int emu_bytes = 0;
    for (int i = 0; i < combined_size - 2; i++) {
        if (combined_payload[i] == 0x00 &&
                combined_payload[i+1] == 0x00 &&
                combined_payload[i+2] <= 0x03) {
            emu_bytes++;
        }
    }
    int rbsp_size = combined_size + emu_bytes;

    // 3. 计算UVLC编码字节数
    int uvlc_bytes = 1;
    int tmp = rbsp_size;
    while (tmp >= 0xFF) {
        uvlc_bytes++;
        tmp -= 0xFF;
    }

    // 4. 分配内存（起始码4 + NAL头1 + payload_type1 + uvlc + rbsp_data）
    int nal_size = 4 + 1 + 1 + uvlc_bytes + rbsp_size;
    uint8_t* sei_nal = (uint8_t*)av_malloc(nal_size + AV_INPUT_BUFFER_PADDING_SIZE);
    if (!sei_nal) {
        av_free(combined_payload);
        return NULL;
    }
    memset(sei_nal + nal_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);

    // 5. 填充起始码和NAL头
    sei_nal[0] = 0x00; sei_nal[1] = 0x00;
    sei_nal[2] = 0x00; sei_nal[3] = 0x01;
    // SEI NAL type
    sei_nal[4] = 0x06;

    // 6. 填充SEI payload
    uint8_t* p = sei_nal + 5;
    *p++ = payload_type;

    // 6.1 UVLC编码的payload_size
    tmp = rbsp_size;
    while (tmp >= 0xFF) {
        *p++ = 0xFF;
        tmp -= 0xFF;
    }
    *p++ = tmp;

    // 6.2 填充RBSP数据（处理防竞争字节）
    for (int i = 0; i < combined_size; ) {
        if (i < combined_size - 2 &&
                combined_payload[i] == 0x00 &&
                combined_payload[i+1] == 0x00 &&
                combined_payload[i+2] <= 0x03) {
            *p++ = 0x00; *p++ = 0x00; *p++ = 0x03;
            i += 2;
        } else {
            *p++ = combined_payload[i++];
        }
    }
    av_free(combined_payload);

    // 7. 验证大小
    if ((p - sei_nal) != nal_size) {
        av_log(NULL, AV_LOG_ERROR, "Size mismatch: %td vs %d\n",
               p - sei_nal, nal_size);
        av_free(sei_nal);
        return NULL;
    }

    *sei_nalu_size = nal_size;
    return sei_nal;
}

bool RtmpPushStream::init()
{
    h264_codec_ctx = init_encoder(m_image_width, m_image_height);
    if (!h264_codec_ctx) {
        qDebug() << "Encoder init failed";
        return false;
    }

    h264_ofmt_ctx = init_rtmp_output(m_rtmpAddr.toStdString().c_str(), h264_codec_ctx);
    if (!h264_ofmt_ctx) {
        // // 释放已分配的编码器
        avcodec_free_context(&h264_codec_ctx);
        return false;
    }

    m_bInitSuccess = true;

    return true;
}



void RtmpPushStream::run()
{

    // 2. 初始化静态变量（帧率控制）
    static const int TARGET_FPS = 30;
    static const int64_t FRAME_INTERVAL_US = 1000000 / TARGET_FPS; // 微秒
    static QElapsedTimer frameTimer;
    static int64_t lastPts = 0;
    static bool isFirstFrame = true;

    while(m_bThreadRun)
    {
        if(m_bInitSuccess)
        {
            if (!h264_ofmt_ctx->pb || !h264_ofmt_ctx->streams[0]) {
                qDebug() << "Output format not initialized!";
            }


            QMutexLocker lk(&m_packetMutex);
            bool b  = m_packetList.isEmpty();
            if(!b)
            {
                AVPacket* packet = m_packetList.takeFirst();
                lk.unlock();

                // 1. 检查数据包有效性
                if (packet->size <= 0 || !packet->data) {
                    qDebug() << "Invalid packet";
                    av_packet_unref(packet);
                    continue;
                }

                if (isFirstFrame) {
                    frameTimer.start();
                    isFirstFrame = false;
                }

                // 3. 设置时间戳（严格按30fps计算）
                AVRational time_base = {1, TARGET_FPS};
                packet->stream_index = h264_ofmt_ctx->streams[0]->index;
                packet->pts = lastPts;
                packet->dts = packet->pts;
                // 每帧递增1（time_base=1/30）
                lastPts += time_base.den;

                // 检测关键帧（通过 NAL 单元类型）
                if (packet->size >= 4) {
                    // // H.264 NAL 单元类型
                    uint8_t nal_type = packet->data[4] & 0x1F;
                    // 7:SPS, 5:IDR帧
                    if (nal_type == 7 || nal_type == 5) {
                        packet->flags |= AV_PKT_FLAG_KEY;
                        // 分割NALU单元，一个AVPacket裸流会存在多个NALU单元
                        int idrIndex = splitNalus(packet->data, packet->size);
                        // 插入SEI信息
                        if(idrIndex > 0 && idrIndex < packet->size)
                        {
                            int sei_size = 0;
                            // 假设SEI段数据长度为345
                            uint8_t sei_payload[345] = {0};
                            for(int i = 0; i < 345; ++i){
                                sei_payload[i] = (i+1)% 255;
                            }

                            uint8_t * sei_nalu = create_sei_nal_unit((char*)sei_payload, 345, &sei_size);
                            if(sei_nalu)
                            {
                                // 2. 创建新缓冲区（深拷贝）
                                int new_size = idrIndex + sei_size + (packet->size - idrIndex);
                                uint8_t* new_data = (uint8_t*)av_malloc(new_size + AV_INPUT_BUFFER_PADDING_SIZE);

                                // 3. 组装数据
                                // 原数据前半部分
                                memcpy(new_data, packet->data, idrIndex);
                                // SEI数据
                                memcpy(new_data + idrIndex, sei_nalu, sei_size);
                                //qDebug() << "sei data size: " << sei_size << QByteArray((char*)sei_nalu, sei_size).toHex(' ');

                                // 原数据后半部分
                                memcpy(new_data + idrIndex + sei_size, packet->data + idrIndex, packet->size - idrIndex);

                                // 4. 替换AVPacket数据（自动释放旧数据）
                                av_packet_from_data(packet, new_data, new_size);

                                // 5. 释放临时SEI NALU
                                av_free(sei_nalu);
                            }
                        }
                    }
                }

                //qDebug() << packet->size << packet->stream_index << packet->pts << packet->dts;

                int ret = av_interleaved_write_frame(h264_ofmt_ctx, packet);
                if(ret < 0)
                {
                    char err_buf[128]={0};
                    av_make_error_string(err_buf, 128, ret);
                    qDebug() << "av_interleaved_write_frame error: " << ret << err_buf;
                }

                // 6. 精确帧率控制（动态补偿）
                qint64 elapsedUs = frameTimer.nsecsElapsed() / 1000;
                qint64 remainingUs = FRAME_INTERVAL_US - elapsedUs;
                if (remainingUs > 0) {
                    QThread::usleep(remainingUs);
                }
                frameTimer.restart();

                av_packet_unref(packet);
            }

        }
        else
        {
            QThread::msleep(3);
        }
    }
}
